#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on 20250417
@author: wangyk
"""

import argparse
import sys
from pathlib import Path
import subprocess
import logging

# 配置logging
logging.basicConfig(level=logging.INFO,
                    format='[%(levelname)s]-[%(asctime)s]-%(message)s')


# 改为相对路径或从环境变量读取
SCRIPT_DIR = Path(
    "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/R/Convertional_Transcriptome/Molecular_Docking_cmd/")

PREDICT_CENTER_COORDINATE = SCRIPT_DIR / "v1.1/0.predict_center_coordinate/1.src.py"
AUTO_BOX = SCRIPT_DIR / "v1.1/1.autobox_py/1.src.py"
PDB_2_PDBQT = SCRIPT_DIR / "v1.1/2.pdb2pdbqt/1.src.py"
DOCKING = SCRIPT_DIR / "docking"


def run_command(cmd, description=""):
    """安全执行命令并捕获输出"""
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"错误: {description}失败", file=sys.stderr)
        print(f"STDERR: {result.stderr}", file=sys.stderr)
        sys.exit(result.returncode)
    return result


def main(args):
    try:
        ligand_path = Path(args.ligand_pdb).resolve()
        output_dir = Path(args.output_dir).resolve()

        # 验证输入
        if not ligand_path.exists():
            raise FileNotFoundError(f"配体文件不存在: {ligand_path}")
        output_dir.mkdir(parents=True, exist_ok=True)

        # 步骤1: 预测中心坐标
        logging.info("步骤1: 预测中心坐标")
        run_command(
            [str(PREDICT_CENTER_COORDINATE), str(ligand_path)],
            "预测中心坐标"
        )

        # 步骤2: 自动计算盒子
        logging.info("步骤2: 自动计算盒子")
        prediction_file = f"{ligand_path}_predictions.csv"
        run_command(
            [str(AUTO_BOX), str(ligand_path), '-p', str(prediction_file)],
            "计算对接盒子"
        )

        # 步骤3: 转换PDBQT
        logging.info("步骤3: PDB转PDBQT")
        json_file = f"{ligand_path}_docking_center_and_autobox_info.json"
        run_command(
            [str(PDB_2_PDBQT), str(ligand_path), str(json_file)],
            "PDB转PDBQT"
        )

        # 步骤4: 执行对接
        logging.info("步骤4: 分子对接")
        prepared_pdbqt = ligand_path.with_name(
            f"{ligand_path.stem}_prepared.pdbqt")
        box_file = prepared_pdbqt.with_suffix('.box.txt')
        run_command(
            [
                str(DOCKING),
                "-p", str(prepared_pdbqt),
                "-c", str(box_file),
                "-od", str(output_dir)
            ],
            "分子对接"
        )

        logging.info("分析完成")

    except Exception as e:
        print(f"其他错误: {str(e)}", file=sys.stderr)
        logging.info("意外退出")
        sys.exit(1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="分子对接，通过prank预测配体中心点，利用PDBParser计算对接box大小，通过mk_prepare_receptor.py转换PDBQT，最后执行分子对接。",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("ligand_pdb", type=str, help="输入配体PDB文件路径。需要去水加氢以及去除其他非配体小分子后的PDB文件。")
    parser.add_argument("-od", "--output_dir", type=str, default="./",metavar='',
                        help="输出目录路径")
    args = parser.parse_args()
    main(args)
